{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp utils\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utils\n",
    "\n",
    "Utils functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import os\n",
    "import pickle\n",
    "import re\n",
    "from typing import Union\n",
    "from inspect import getmembers\n",
    "from collections import defaultdict\n",
    "\n",
    "from loguru import logger\n",
    "import logging\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import transformers\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from sklearn.preprocessing import MultiLabelBinarizer\n",
    "from transformers import AutoTokenizer, PreTrainedTokenizer, TFPreTrainedModel\n",
    "from m3tl.special_tokens import TRAIN, EVAL, PREDICT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-05 23:39:46.452 | INFO     | m3tl.base_params:register_multiple_problems:476 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-05 23:39:46.453 | INFO     | m3tl.base_params:register_multiple_problems:476 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-05 23:39:46.453 | INFO     | m3tl.base_params:register_multiple_problems:476 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-05 23:39:46.454 | INFO     | m3tl.base_params:register_multiple_problems:476 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-05 23:39:46.454 | INFO     | m3tl.base_params:register_multiple_problems:476 - Adding new problem weibo_fake_regression, problem type: regression\n",
      "2021-06-05 23:39:46.455 | INFO     | m3tl.base_params:register_multiple_problems:476 - Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "2021-06-05 23:39:46.455 | INFO     | m3tl.base_params:register_multiple_problems:476 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-05 23:39:46.456 | WARNING  | m3tl.base_params:prepare_dir:314 - bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from m3tl.test_base import TestBase\n",
    "import m3tl\n",
    "import shutil\n",
    "import numpy as np\n",
    "test_base = TestBase()\n",
    "params = test_base.params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def load_transformer_tokenizer(tokenizer_name: str, load_module_name=None):\n",
    "    \"\"\"some tokenizers cannot be loaded using AutoTokenizer.\n",
    "\n",
    "    this function served as a util function to catch that situation.\n",
    "\n",
    "    Args:\n",
    "        tokenizer_name (str): tokenizer name\n",
    "    \"\"\"\n",
    "    if load_module_name:\n",
    "        tok = getattr(transformers, load_module_name).from_pretrained(\n",
    "            tokenizer_name)\n",
    "    else:\n",
    "        tok = AutoTokenizer.from_pretrained(tokenizer_name)\n",
    "\n",
    "    return tok"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PreTrainedTokenizer(name_or_path='voidful/albert_chinese_tiny', vocab_size=21128, model_max_len=1000000000000000019884624838656, is_fast=False, padding_side='right', special_tokens={'unk_token': '[UNK]', 'sep_token': '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token': '[MASK]'})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "load_transformer_tokenizer(\n",
    "            'voidful/albert_chinese_tiny', 'BertTokenizer')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def load_transformer_config(config_name_or_dict, load_module_name=None):\n",
    "    \"\"\"Some models need specify loading module\n",
    "\n",
    "    Args:\n",
    "        config_name (str): module name\n",
    "        load_module_name (str, optional): loading module name. Defaults to None.\n",
    "\n",
    "    Returns:\n",
    "        config: config\n",
    "    \"\"\"\n",
    "    if load_module_name:\n",
    "        load_module = getattr(transformers, load_module_name)\n",
    "    else:\n",
    "        load_module = transformers.AutoConfig\n",
    "    if isinstance(config_name_or_dict, str):\n",
    "        config = load_module.from_pretrained(\n",
    "            config_name_or_dict, output_attentions=True, output_hidden_states=True)\n",
    "    elif isinstance(config_name_or_dict, dict):\n",
    "        config = load_module.from_dict(\n",
    "            config_name_or_dict, output_attentions=True, output_hidden_states=True)\n",
    "    else:\n",
    "        raise ValueError('config_name_or_dict should be str or dict')\n",
    "    return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load config with name\n",
    "config = load_transformer_config(\n",
    "    'bert-base-chinese')\n",
    "config_dict = config.to_dict()\n",
    "# load config with dict\n",
    "config = load_transformer_config(\n",
    "    config_dict, load_module_name='BertConfig')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def load_transformer_model(model_name_or_config, load_module_name=None):\n",
    "    if load_module_name:\n",
    "        load_module = getattr(transformers, load_module_name)\n",
    "    else:\n",
    "        load_module = transformers.TFAutoModel\n",
    "\n",
    "    if isinstance(model_name_or_config, str):\n",
    "        try:\n",
    "            model = load_module.from_pretrained(\n",
    "                model_name_or_config, output_attentions=True, output_hidden_states=True)\n",
    "        except OSError:\n",
    "            model = load_module.from_pretrained(\n",
    "                model_name_or_config, from_pt=True, output_attentions=True, output_hidden_states=True)\n",
    "    else:\n",
    "        model = load_module(model_name_or_config)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.decoder.weight', 'predictions.dense.bias', 'predictions.bias', 'predictions.LayerNorm.bias', 'predictions.dense.weight', 'predictions.decoder.bias', 'predictions.LayerNorm.weight']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n"
     ]
    }
   ],
   "source": [
    "# load by name(load weights)\n",
    "# this is a pt only model\n",
    "model = load_transformer_model(\n",
    "    'voidful/albert_chinese_tiny')\n",
    "\n",
    "# load by config (not load weights)\n",
    "model = load_transformer_model(load_transformer_config(\n",
    "    'bert-base-chinese'), 'TFBertModel')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "def get_label_encoder_save_path(params, problem: str) -> str:\n",
    "    if not hasattr(params, 'ckpt_dir'):\n",
    "        raise ValueError('Seems no problem assigned to params. Please check.')\n",
    "    problem_path = params.ckpt_dir\n",
    "    create_path(problem_path)\n",
    "    problem_type = params.problem_type[problem]\n",
    "    le_path = os.path.join(problem_path, '%s_label_encoder.pkl' % problem)\n",
    "    return le_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LabelEncoder(BaseEstimator, TransformerMixin):\n",
    "\n",
    "    def fit(self, y):\n",
    "        \"\"\"Fit label encoder\n",
    "        Parameters\n",
    "        ----------\n",
    "        y : array-like of shape (n_samples,)\n",
    "            Target values.\n",
    "        Returns\n",
    "        -------\n",
    "        self : returns an instance of self.\n",
    "        \"\"\"\n",
    "        self.encode_dict = {}\n",
    "        self.decode_dict = {}\n",
    "        label_set = set(y)\n",
    "\n",
    "        label_set = sorted(list(label_set))\n",
    "\n",
    "        for l_ind, l in enumerate(label_set):\n",
    "\n",
    "            new_ind = l_ind\n",
    "\n",
    "            self.encode_dict[l] = new_ind\n",
    "            self.decode_dict[new_ind] = l\n",
    "\n",
    "        return self\n",
    "\n",
    "    def fit_transform(self, y):\n",
    "        \"\"\"Fit label encoder and return encoded labels\n",
    "        Parameters\n",
    "        ----------\n",
    "        y : array-like of shape [n_samples]\n",
    "            Target values.\n",
    "        Returns\n",
    "        -------\n",
    "        y : array-like of shape [n_samples]\n",
    "        \"\"\"\n",
    "        self.fit(y)\n",
    "        y = self.transform(y)\n",
    "        return y\n",
    "\n",
    "    def transform(self, y):\n",
    "        \"\"\"Transform labels to normalized encoding.\n",
    "        Parameters\n",
    "        ----------\n",
    "        y : array-like of shape [n_samples]\n",
    "            Target values.\n",
    "        Returns\n",
    "        -------\n",
    "        y : array-like of shape [n_samples]\n",
    "        \"\"\"\n",
    "        encode_y = []\n",
    "        for l in y:\n",
    "            encode_y.append(self.encode_dict[l])\n",
    "\n",
    "        return np.array(encode_y)\n",
    "\n",
    "    def inverse_transform(self, y):\n",
    "        \"\"\"Transform labels back to original encoding.\n",
    "        Parameters\n",
    "        ----------\n",
    "        y : numpy array of shape [n_samples]\n",
    "            Target values.\n",
    "        Returns\n",
    "        -------\n",
    "        y : numpy array of shape [n_samples]\n",
    "        \"\"\"\n",
    "        decode_y = []\n",
    "        for l in y:\n",
    "            decode_y.append(self.decode_dict[l])\n",
    "\n",
    "        return np.array(decode_y)\n",
    "\n",
    "    def dump(self, path):\n",
    "        with open(path, 'wb') as f:\n",
    "            pickle.dump(self.decode_dict, f)\n",
    "\n",
    "    def load(self, path):\n",
    "        with open(path, 'rb') as f:\n",
    "            self.decode_dict = pickle.load(f)\n",
    "\n",
    "        self.encode_dict = {v: k for k, v in self.decode_dict.items()}\n",
    "\n",
    "\n",
    "def create_path(path):\n",
    "    if not os.path.exists(path):\n",
    "        os.makedirs(path, exist_ok=True)\n",
    "\n",
    "def need_make_label_encoder(mode: str, le_path: str, overwrite=False):\n",
    "    return mode == TRAIN and (not os.path.exists(le_path) or overwrite)\n",
    "\n",
    "def get_or_make_label_encoder(params, problem: str, mode: str, label_list=None, overwrite=True) -> Union[LabelEncoder, MultiLabelBinarizer, PreTrainedTokenizer]:\n",
    "    \"\"\"Function to unify ways to get or create label encoder for various\n",
    "    problem type.\n",
    "\n",
    "    cls: LabelEncoder\n",
    "    seq_tag: LabelEncoder\n",
    "    multi_cls: MultiLabelBinarizer\n",
    "    seq2seq_text: Tokenizer\n",
    "\n",
    "    Arguments:\n",
    "        problem {str} -- problem name\n",
    "        mode {mode} -- mode\n",
    "\n",
    "    Keyword Arguments:\n",
    "        label_list {list} -- label list to fit the encoder (default: {None})\n",
    "\n",
    "    Returns:\n",
    "        LabelEncoder -- label encoder\n",
    "    \"\"\"\n",
    "\n",
    "    problem_path = params.ckpt_dir\n",
    "    create_path(problem_path)\n",
    "    problem_type = params.problem_type[problem]\n",
    "\n",
    "    get_or_make_custom_le_fn = params.get_or_make_label_encoder_fn_dict[\n",
    "        problem_type]\n",
    "    label_encoder = get_or_make_custom_le_fn(params, problem, mode, label_list, overwrite=overwrite)\n",
    "\n",
    "    params.set_problem_info(problem=problem, info_name='label_encoder_called', info=True)\n",
    "    return label_encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "le_train = get_or_make_label_encoder(\n",
    "    params=params, problem='weibo_fake_ner', mode=m3tl.TRAIN, label_list=[['a', 'b'], ['c']]\n",
    ")\n",
    "# seq_tag will add [PAD]\n",
    "assert len(le_train.encode_dict) == 4, le_train.encode_dict\n",
    "\n",
    "le_predict = get_or_make_label_encoder(\n",
    "    params=params, problem='weibo_fake_ner', mode=m3tl.PREDICT)\n",
    "assert le_predict.encode_dict==le_train.encode_dict\n",
    "\n",
    "# list train\n",
    "le_train = get_or_make_label_encoder(\n",
    "    params=params, problem='weibo_fake_cls', mode=m3tl.TRAIN, label_list=['a', 'b', 'c']\n",
    ")\n",
    "# seq_tag will add [PAD]\n",
    "assert len(le_train.encode_dict) == 3\n",
    "\n",
    "le_predict = get_or_make_label_encoder(\n",
    "    params=params, problem='weibo_fake_cls', mode=m3tl.PREDICT)\n",
    "assert le_predict.encode_dict==le_train.encode_dict\n",
    "\n",
    "# text\n",
    "le_train = get_or_make_label_encoder(\n",
    "    params=params, problem='weibo_masklm', mode=m3tl.TRAIN)\n",
    "assert isinstance(le_train, transformers.PreTrainedTokenizer)\n",
    "le_predict = get_or_make_label_encoder(\n",
    "    params=params, problem='weibo_masklm', mode=m3tl.PREDICT)\n",
    "assert isinstance(le_predict, transformers.PreTrainedTokenizer)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def cluster_alphnum(text: str) -> list:\n",
    "    \"\"\"Simple funtions to aggregate eng and number\n",
    "\n",
    "    Arguments:\n",
    "        text {str} -- input text\n",
    "\n",
    "    Returns:\n",
    "        list -- list of string with chinese char or eng word as element\n",
    "    \"\"\"\n",
    "    return_list = []\n",
    "    last_is_alphnum = False\n",
    "\n",
    "    for char in text:\n",
    "        is_alphnum = bool(re.match('^[a-zA-Z0-9\\[]+$', char))\n",
    "        is_right_brack = char == ']'\n",
    "        if is_alphnum:\n",
    "            if last_is_alphnum:\n",
    "                return_list[-1] += char\n",
    "            else:\n",
    "                return_list.append(char)\n",
    "                last_is_alphnum = True\n",
    "        elif is_right_brack:\n",
    "            if return_list:\n",
    "                return_list[-1] += char\n",
    "            else:\n",
    "                return_list.append(char)\n",
    "            last_is_alphnum = False\n",
    "        else:\n",
    "            return_list.append(char)\n",
    "            last_is_alphnum = False\n",
    "    return return_list\n",
    "\n",
    "\n",
    "def filter_empty(input_list, target_list):\n",
    "    \"\"\"Filter empty inputs or targets\n",
    "\n",
    "    Arguments:\n",
    "        input_list {list} -- input list\n",
    "        target_list {list} -- target list\n",
    "\n",
    "    Returns:\n",
    "        input_list, target_list -- data after filter\n",
    "    \"\"\"\n",
    "    return_input, return_target = [], []\n",
    "    for inp, tar in zip(input_list, target_list):\n",
    "        if inp and tar:\n",
    "            return_input.append(inp)\n",
    "            return_target.append(tar)\n",
    "    return return_input, return_target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "\n",
    "def infer_shape_and_type_from_dict(inp_dict: dict, fix_dim_for_high_rank_tensor=True):\n",
    "    shape_dict = {}\n",
    "    type_dict = {}\n",
    "    for feature_name, feature in inp_dict.items():\n",
    "        if type(feature) is list:\n",
    "            feature = np.array(feature)\n",
    "        if type(feature) is np.ndarray:\n",
    "            if np.issubdtype(feature.dtype, np.integer):\n",
    "                type_dict[feature_name] = tf.int32\n",
    "            elif np.issubdtype(feature.dtype, np.floating):\n",
    "                type_dict[feature_name] = tf.float32\n",
    "\n",
    "            # this seems not a good idea\n",
    "            if len(feature.shape) > 1 and fix_dim_for_high_rank_tensor:\n",
    "                shape_dict[feature_name] = [\n",
    "                    None] + list(feature.shape[1:])\n",
    "            else:\n",
    "                shape_dict[feature_name] = [\n",
    "                    None for _ in feature.shape]\n",
    "\n",
    "        elif np.issubdtype(type(feature), np.floating):\n",
    "\n",
    "            type_dict[feature_name] = tf.float32\n",
    "            shape_dict[feature_name] = []\n",
    "        elif np.issubdtype(type(feature), np.integer):\n",
    "\n",
    "            type_dict[feature_name] = tf.int32\n",
    "            shape_dict[feature_name] = []\n",
    "        else:\n",
    "            if isinstance(feature, str):\n",
    "                feature = feature.encode('utf8')\n",
    "\n",
    "            type_dict[feature_name] = tf.string\n",
    "            shape_dict[feature_name] = []\n",
    "    return shape_dict, type_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dose not support nested dict\n",
    "test_dict = {\n",
    "    'test1': np.random.uniform(size=(64, 32)),\n",
    "    'test2': np.array([1, 2, 3], dtype='int32'),\n",
    "    'test5': 5\n",
    "}\n",
    "desc_dict = infer_shape_and_type_from_dict(\n",
    "    test_dict)\n",
    "assert desc_dict == ({'test1': [None, 32], 'test2': [None], 'test5': []}, {\n",
    "                    'test1': tf.float32, 'test2': tf.int32, 'test5': tf.int32})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_transformer_main_model(model, key='embeddings'):\n",
    "    \"\"\"Function to extrac model name from huggingface transformer models.\n",
    "\n",
    "    Args:\n",
    "        model (Model): Huggingface transformers model\n",
    "        key (str, optional): Key to identify model. Defaults to 'embeddings'.\n",
    "\n",
    "    Returns:\n",
    "        model\n",
    "    \"\"\"\n",
    "\n",
    "    for attr_name, attr in getmembers(model):\n",
    "        if attr_name == key:\n",
    "            return model\n",
    "        if hasattr(attr, key):\n",
    "            return attr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.decoder.weight', 'predictions.dense.bias', 'predictions.bias', 'predictions.LayerNorm.bias', 'predictions.dense.weight', 'predictions.decoder.bias', 'predictions.LayerNorm.weight']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = load_transformer_model(\n",
    "    'voidful/albert_chinese_tiny')\n",
    "main_model = get_transformer_main_model(model)\n",
    "isinstance(main_model, transformers.TFAlbertMainLayer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "def get_embedding_table_from_model(model: TFPreTrainedModel):\n",
    "    embedding_layer = model.get_input_embeddings()\n",
    "    return embedding_layer.weights[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding = get_embedding_table_from_model(\n",
    "    model)\n",
    "assert embedding.shape == (21128, 128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_shape_list(tensor, expected_rank=None, name=None):\n",
    "    \"\"\"Returns a list of the shape of tensor, preferring static dimensions.\n",
    "\n",
    "    Args:\n",
    "      tensor: A tf.Tensor object to find the shape of.\n",
    "      expected_rank: (optional) int. The expected rank of `tensor`. If this is\n",
    "        specified and the `tensor` has a different rank, and exception will be\n",
    "        thrown.\n",
    "      name: Optional name of the tensor for the error message.\n",
    "\n",
    "    Returns:\n",
    "      A list of dimensions of the shape of tensor. All static dimensions will\n",
    "      be returned as python integers, and dynamic dimensions will be returned\n",
    "      as tf.Tensor scalars.\n",
    "    \"\"\"\n",
    "    shape = tensor.shape.as_list()\n",
    "\n",
    "    non_static_indexes = []\n",
    "    for (index, dim) in enumerate(shape):\n",
    "        if dim is None:\n",
    "            non_static_indexes.append(index)\n",
    "\n",
    "    if not non_static_indexes:\n",
    "        return shape\n",
    "\n",
    "    dyn_shape = tf.shape(input=tensor)\n",
    "    for index in non_static_indexes:\n",
    "        shape[index] = dyn_shape[index]\n",
    "    return shape\n",
    "\n",
    "\n",
    "def gather_indexes(sequence_tensor, positions):\n",
    "    \"\"\"Gathers the vectors at the specific positions over a minibatch.\"\"\"\n",
    "    sequence_shape = get_shape_list(sequence_tensor, expected_rank=3)\n",
    "    batch_size = sequence_shape[0]\n",
    "    seq_length = sequence_shape[1]\n",
    "    width = sequence_shape[2]\n",
    "\n",
    "    flat_offsets = tf.reshape(\n",
    "        tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])\n",
    "    flat_offsets = tf.cast(flat_offsets, tf.int64)\n",
    "    flat_positions = tf.reshape(positions + flat_offsets, [-1])\n",
    "    flat_sequence_tensor = tf.reshape(sequence_tensor,\n",
    "                                      [batch_size * seq_length, width])\n",
    "    # https://github.com/tensorflow/tensorflow/issues/36236\n",
    "    output_tensor = tf.gather(flat_sequence_tensor*1, flat_positions)\n",
    "    return output_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "def dispatch_features(features, hidden_feature, problem, mode):\n",
    "    # get features with ind == 1\n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        feature_this_round = features\n",
    "        hidden_feature_this_round = hidden_feature\n",
    "    else:\n",
    "        multiplier_name = '%s_loss_multiplier' % problem\n",
    "\n",
    "        # record_ind = tf.where(tf.cast(\n",
    "        #     tf.squeeze(features[multiplier_name]), tf.bool))\n",
    "\n",
    "        record_ind = tf.where(tf.cast(features[multiplier_name], tf.bool))\n",
    "\n",
    "        hidden_feature_this_round = {}\n",
    "        for hidden_feature_name in hidden_feature:\n",
    "            if hidden_feature_name != 'embed_table':\n",
    "                hidden_feature_this_round[hidden_feature_name] = tf.squeeze(tf.gather(\n",
    "                    hidden_feature[hidden_feature_name], record_ind, axis=0\n",
    "                ), axis=1)\n",
    "                hidden_feature_this_round[hidden_feature_name].set_shape(\n",
    "                    [None, *hidden_feature[hidden_feature_name].shape.as_list()[1:]])\n",
    "            else:\n",
    "                hidden_feature_this_round[hidden_feature_name] = hidden_feature[hidden_feature_name]\n",
    "\n",
    "        feature_this_round = {}\n",
    "        for features_name in features:\n",
    "            feature_this_round[features_name] = tf.gather_nd(\n",
    "                features[features_name],\n",
    "                record_ind)\n",
    "\n",
    "    return feature_this_round, hidden_feature_this_round"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def _add_to_dict(loss_dict, model, added_name, ele_name):\n",
    "    ele_list = getattr(model, ele_name)\n",
    "    for ele in ele_list:\n",
    "        try:\n",
    "            ele_tf_name = ele.name\n",
    "            if ele_tf_name not in added_name:\n",
    "                loss_dict[ele_name].append(ele)\n",
    "                added_name.append(ele_tf_name)\n",
    "        except AttributeError:\n",
    "            if ele not in added_name:\n",
    "                loss_dict[ele_name].append(ele)\n",
    "                added_name.append(ele)\n",
    "    return loss_dict, added_name\n",
    "\n",
    "def create_dict_from_nested_model(model: tf.keras.Model, loss_dict=None, ele_name='losses', added_name=None) -> dict:\n",
    "    if loss_dict is None:\n",
    "        loss_dict = defaultdict(list)\n",
    "\n",
    "    if added_name is None:\n",
    "        added_name = []\n",
    "    if not isinstance(model, tf.keras.Model):\n",
    "        loss_dict, added_name = _add_to_dict(loss_dict, model, added_name, ele_name)\n",
    "\n",
    "    else:\n",
    "        for layer in model.layers:\n",
    "            layer_loss_dict = loss_dict.get(layer.name, None)\n",
    "            layer_loss_dict = create_dict_from_nested_model(layer, layer_loss_dict, ele_name, added_name)\n",
    "            loss_dict[layer.name] = layer_loss_dict\n",
    "        loss_dict, added_name = _add_to_dict(loss_dict, model, added_name, ele_name)\n",
    "\n",
    "    # remove empty\n",
    "    loss_dict_keys = list(loss_dict.keys())\n",
    "    for k in loss_dict_keys:\n",
    "        if not loss_dict[k]:\n",
    "            del loss_dict[k]\n",
    "    return loss_dict\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def variable_summaries(var, name):\n",
    "    \"\"\"Attach a lot of summaries to a Tensor (for TensorBoard visualization).\"\"\"\n",
    "    with tf.name_scope(name):\n",
    "        mean = tf.reduce_mean(input_tensor=var)\n",
    "        tf.summary.scalar('mean', mean)\n",
    "        with tf.name_scope('stddev'):\n",
    "            stddev = tf.sqrt(tf.reduce_mean(\n",
    "                input_tensor=tf.square(var - mean)))\n",
    "        tf.summary.scalar('stddev', stddev)\n",
    "        tf.summary.scalar('max', tf.reduce_max(input_tensor=var))\n",
    "        tf.summary.scalar('min', tf.reduce_min(input_tensor=var))\n",
    "        tf.summary.histogram('histogram', var)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "M3TL_PHASE = TRAIN\n",
    "def set_phase(phase: str):\n",
    "    if phase not in [TRAIN, EVAL, PREDICT]:\n",
    "        raise ValueError('expect one of following values: {}, got: {}'.format([TRAIN, EVAL, PREDICT], phase))\n",
    "    logger.info('Setting phase to {}'.format(phase))\n",
    "    global M3TL_PHASE\n",
    "    M3TL_PHASE = phase\n",
    "\n",
    "def get_phase() -> str:\n",
    "    return M3TL_PHASE\n",
    "\n",
    "IS_PYSPARK = False\n",
    "def set_is_pyspark(is_pyspark: bool):\n",
    "    global IS_PYSPARK\n",
    "    IS_PYSPARK = is_pyspark\n",
    "\n",
    "def get_is_pyspark() -> bool:\n",
    "    return IS_PYSPARK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "class TFRedundantWarningFilter(logging.Filter):\n",
    "    def filter(self, record: logging.LogRecord) -> bool:\n",
    "        msg = record.msg\n",
    "        lvl = record.levelname\n",
    "        if lvl in ('WARNING', 'U'):\n",
    "            if \".optimizer's state\" in msg:\n",
    "                return False\n",
    "\n",
    "            if \"AutoGraph could not\" in msg:\n",
    "                return False \n",
    "\n",
    "            if \"Converting sparse IndexedSlices\" in msg:\n",
    "                return False\n",
    "        return True\n",
    "\n",
    "def compress_tf_warnings():\n",
    "    warnings.simplefilter('ignore')\n",
    "    tf.get_logger().addFilter(TFRedundantWarningFilter())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
